#!/usr/bin/env python3
'''
未验证

训练记录：
20241002：训练过程中的平局奖励很快到达了9000+，但是在保存最好模型的时候，发现最好模型的奖励只有5000+，上升缓慢，确保保存模型时的代码正确性
20241003：发现由于ga算法是遗传、采样的方式生成的网络，所以即使是发现了最好的采样方式，再次采样生成不一定是最好的网络，所以在测试保存的时候，需要对前N个最好的
采样进行多次测试，选择最好的网络进行保存，增加再次采样得到最好回报的网络的概率，目前通过这个方法已经完成了训练

待做
- [ ] 学习以下ga算法的原理
'''
import gymnasium as gym
import collections
import copy
import time
import numpy as np
import os

import torch
import torch.nn as nn
import torch.multiprocessing as mp

from tensorboardX import SummaryWriter


NOISE_STD = 0.01
POPULATION_SIZE = 2000 # 种群的数量，网络的数量
PARENTS_COUNT = 10 # todo 要进行遗传的父母个数？
WORKERS_COUNT = 6 # 子进程的数量
SEEDS_PER_WORKER = POPULATION_SIZE // WORKERS_COUNT # todo 每个子进程需要的SEEDS
MAX_SEED = 2**32 - 1
TEST_FREQUENCY = 15


class Net(nn.Module):
    def __init__(self, obs_size, act_size, hid_size=64):
        super(Net, self).__init__()

        self.mu = nn.Sequential(
            nn.Linear(obs_size, hid_size),
            nn.Tanh(),
            nn.Linear(hid_size, hid_size),
            nn.Tanh(),
            nn.Linear(hid_size, act_size),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.mu(x)


def evaluate(env, net):
    '''
    评估网路
    '''
    obs, _ = env.reset()
    reward = 0.0
    steps = 0
    while True:
        obs_v = torch.FloatTensor(np.array([obs]))
        action_v = net(obs_v)
        obs, r, done, trunc, _ = env.step(action_v.data.numpy()[0])
        reward += r
        steps += 1
        if done or trunc:
            break
    return reward, steps


def mutate_net(net, seed, copy_net=True):
    # 拷贝网络
    new_net = copy.deepcopy(net) if copy_net else net
    # 设置随机数
    np.random.seed(seed)
    for p in new_net.parameters():
        # 生成和网路层参数相同维度的噪声
        noise_t = torch.tensor(np.random.normal(size=p.data.size()).astype(np.float32))
        p.data += NOISE_STD * noise_t
    return new_net


def build_net(env, seeds):
    '''
    构建网络
    '''
    # 设置种子，设置完成后，构建Net后，其内部的种子就是设置的种子了
    # 为什么这边的种子仅需要设置torch.manual_seed(seeds[0])
    #   而其余的种子设置到mutate_net里面
    # 因为seeds中记录的是种群遗传算法遗传链，第一个种子是初始网络的种子
    # 后续种子都是基于前一个网络的种子进行变异得到的
    # 所以在构建Net的时候，只需要设置第一个种子即可
    # 而在变异的时候，需要将后续所有的种子都设置到网络中
    torch.manual_seed(seeds[0])
    net = Net(env.observation_space.shape[0], env.action_space.shape[0])
    for seed in seeds[1:]:
        # 网络变异，会把除了初始seed都拿去变异一次
        net = mutate_net(net, seed, copy_net=False)
    return net


OutputItem = collections.namedtuple('OutputItem', field_names=['seeds', 'reward', 'steps'])


def worker_func(input_queue, output_queue):
    '''
    每个工作进程有独立的环境
    '''
    env = gym.make("InvertedDoublePendulum-v4")
    # todo 貌似整个代码并未限制net_seeds长度，那么是否会无限制增长？
    cache = {}

    while True:
        # 得到父母网络的种子
        parents = input_queue.get()
        if parents is None:
            break
        new_cache = {}
        for net_seeds in parents:
            if len(net_seeds) > 1:
                # 如果种子的长度大于1，说明是需要进行遗传的父母
                # 根据种子从缓存中得到网络
                net = cache.get(net_seeds[:-1])
                if net is not None:
                    # 如果缓存中存在网络，说明之前的遗传链已经缓存到了当前子进程中
                    # 所以
                    # 构建的时候会将所有的seed全部都设置到网络中
                    # 而更新的时候仅拿最新的seed设置到网络中
                    net = mutate_net(net, net_seeds[-1])
                else:
                    # 如果网络不存在，则说明当前子进程没有缓存该网络
                    # 需要重头构建遗传链，所以使用build
                    net = build_net(env, net_seeds)
            else:
                net = build_net(env, net_seeds)
            # 以网络种子为键，网络为值，放入缓存中
            new_cache[net_seeds] = net
            reward, steps = evaluate(env, net)
            output_queue.put(OutputItem(seeds=net_seeds, reward=reward, steps=steps))
        cache = new_cache

import ptan


def test_net(net, env, count=10, device="cpu"):
    '''
    测试网络，专门用于测试训练的结果，不需要进行探索

    net: 网络
    env: 环境
    count: 测试次数
    device: 设备

    return： （平均奖励，平均步数）
    '''
    rewards = 0.0
    steps = 0
    for _ in range(count):
        # 重置环境
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.float32_preprocessor([obs]).to(device)
            # 计算均值，在高斯分布中，均值是所有的分布中，最有可能的值，所以使用均值作为执行的动作值
            mu_v = net(obs_v)[0]
            action = mu_v.squeeze(dim=0).data.cpu().numpy()
            action = np.expand_dims(action, axis=0)
            obs, reward, done, trunc, _ = env.step(action)
            rewards += reward
            steps += 1
            if done or trunc:
                break
    return rewards / count, steps / count


if __name__ == "__main__":
    mp.set_start_method('spawn')
    writer = SummaryWriter(comment="-mujoco-double-pendulum-ga")
    test_env = gym.make("InvertedDoublePendulum-v4")

    save_path = os.path.join("saves", "ga-double-pendulum")
    os.makedirs(save_path, exist_ok=True)

    seed = int(time.time())
    rng = np.random.default_rng(seed)

    # 子进程输入队列，存放的是给子进程的变异种子
    input_queues = []
    # output_queue 子进程返回给主进程的参数队列
    # 内容：（所使用的种子list，奖励，步数）
    output_queue = mp.Queue(maxsize=WORKERS_COUNT)
    workers = []
    for _ in range(WORKERS_COUNT):
        input_queue = mp.Queue(maxsize=1)
        input_queues.append(input_queue)
        w = mp.Process(target=worker_func, args=(input_queue, output_queue))
        w.start()
        # 根据每个子进程的的输入种子的大小SEED_PER_WORKER，得到每个子进程的种子列表
        # seeds存放的东西是[(seed1，，，),(seed2，，，),...]
        # seeds = [(np.random.randint(MAX_SEED),) for _ in range(SEEDS_PER_WORKER)]
        seeds = [(rng.integers(np.iinfo(np.int32).max),) for _ in range(SEEDS_PER_WORKER)]
        # 将得到的种子列表放入输入队列
        input_queue.put(seeds)

    gen_idx = 0
    elite = None
    best_reward = 0
    while True:
        t_start = time.time()
        batch_steps = 0 # 统计所有子进程的步数
        population = [] # 存放的内容是[(seed1, reward1),(seed2, reward2),...]
        while len(population) < SEEDS_PER_WORKER * WORKERS_COUNT:
            # 收集到的种群数量小于目标大小
            # 则继续从子进程中收集
            out_item = output_queue.get()
            population.append((out_item.seeds, out_item.reward))
            batch_steps += out_item.steps
        if elite is not None:
            # 如果评估最好的网络存在，则将其加入种群
            population.append(elite)
        # 对所有种群进行奖励排序，将奖励回报最高的排在最前面
        population.sort(key=lambda p: p[1], reverse=True)
        # 得到排名前PARENTS_COUNT的种群的奖励平均值、最大值、标准差
        rewards = [p[1] for p in population[:PARENTS_COUNT]]
        reward_mean = np.mean(rewards)
        reward_max = np.max(rewards)
        reward_std = np.std(rewards)
        # 记录到tensorboard中
        writer.add_scalar("reward_mean", reward_mean, gen_idx)
        writer.add_scalar("reward_std", reward_std, gen_idx)
        writer.add_scalar("reward_max", reward_max, gen_idx)
        writer.add_scalar("batch_steps", batch_steps, gen_idx)
        writer.add_scalar("gen_seconds", time.time() - t_start, gen_idx)
        # 计算速度
        speed = batch_steps / (time.time() - t_start)
        writer.add_scalar("speed", speed, gen_idx)
        print("%d: reward_mean=%.2f, reward_max=%.2f, reward_std=%.2f, speed=%.2f f/s" % (
            gen_idx, reward_mean, reward_max, reward_std, speed))

        # 保存最好的种子
        elite = population[0]
        for worker_queue in input_queues:
            # 遍历每个子进程的输入队列
            seeds = []
            for _ in range(SEEDS_PER_WORKER):
                # 循环大小是每个进程需要的种子数
                # 得到随机数，随机数的范围是需要遗传父母的个数之内
                # 其实也就是仅选择奖励回报最高的几个seed进行遗传，而不高的seed不进行遗传
                # 这里也就是ga算法的关键
                parent = rng.integers(PARENTS_COUNT)
                # 得到下一个随机数，随机数的范围是0到32位正数最大值
                next_seed = rng.integers(np.iinfo(np.int32).max)
                # 将list(population[parent][0])元组转list
                # 将父母的种子和新种子合并为一个list后再转换为元组，推进到seeds队列中
                # 将新的种子推进到seeds队列
                seeds.append(tuple(list(population[parent][0]) + [next_seed]))
            # 将得到的新种子推进到子进程的输入队列中
            worker_queue.put(seeds)
        if gen_idx % TEST_FREQUENCY == 0:
            best_net = None
            for i in range(PARENTS_COUNT):
                best_seeds = population[i][0]
                for _ in range(10):
                    currnet_net = build_net(test_env, best_seeds)
                    reward, steps = test_net(currnet_net, test_env, count=50)
                    if reward > best_reward:
                        best_reward = reward
                        best_net = currnet_net
                        print(f"Best net evaluated in generation {gen_idx}: reward={reward:.2f}, steps={steps:.2f}")
            if best_net is not None:
                torch.save(best_net.state_dict(), os.path.join(save_path, f"best_model_gen_{gen_idx}_reward_{best_reward}.pth"))
                print(f"Saved best model at generation {gen_idx}")
        gen_idx += 1

    pass
